/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Contains classes for scatterers, i.e. dust grains.

// $Id$


#ifndef __scatterer__
#define __scatterer__

#include <vector> 
#include "boost/shared_ptr.hpp"
#include "mcrx-types.h"
#include "ray.h"
#include "random.h"
#include "constants.h"
#include "angular.h"
#include "chromatic_policy.h"
#include "mono_poly_abstract.h"
#include "interpolatort.h"
#include "mcrx-debug.h"
#include "mcrx-units.h"

namespace mcrx {
  template <template<class> class, typename> class scatterer;
  T_float HG_phase_function(T_float, T_float);

  BZ_DEFINE_BINARY_FUNC(Fn_hgpf, HG_phase_function);
  BZ_DECLARE_ARRAY_ET_BINARY(HG_phase_function, Fn_hgpf);
  BZ_DECLARE_ARRAY_ET_BINARY_SCALAR(HG_phase_function, Fn_hgpf, T_float);

  template<typename> class hg_returntype;
  template<typename T> class hg_returntype<blitz::Array<T, 1> >;

  template <template<class> class, typename> class simple_HG_dust_grain;
  template <template<class> class, typename> class HG_dust_grain;
  template <template<class> class, typename> class Draine_grain;
}


/** Abstract base class for things that scatter rays. This class
    defines the interface for the various dust grain classes. For
    efficiency reasons it also hard codes the opacity and albedo data,
    since it works together with the dust_model class. */
template <template<class> class chromatic_policy, 
	  typename T_rng_policy = mcrx::global_random> 
class mcrx::scatterer  {
  friend class polychromatic_dust_model_policy<scatterer>;
public:
  /** The chromatic_policy class template would normally be publicly
      inherited from.  However, because it needs to know the type of
      this class, it's tricky and since the chromatic_policy has only
      type definitions we can just access these directly without
      inheritance.  */
  typedef chromatic_policy<scatterer> T_chromatic_policy;
  /// Forward definition from policy class.
  typedef typename T_chromatic_policy::T_lambda T_lambda;
  /// Forward definition from policy class.
  typedef typename T_chromatic_policy::T_scatterer_vector T_scatterer_vector;
  /// Forward definition from policy class.
  typedef typename T_chromatic_policy::T_biaser T_biaser;
  /// The ray contains a lambda type, which depends on the chromatic_policy.
  typedef ray<T_lambda> T_ray;
protected:
  /** The opacity of the scatterer as a function of wavelength. This
      is kept in the base class because it interacts with the
      dust_model class. */
  T_lambda opacity_;
  /// The albedo of the scatterer as a function of wavelength.
  T_lambda albedo_;

  /// Default constructor is intended for derived classes.
  scatterer () {};

  /** The units of the opacity curve. (See the comment about the
      units_ member of the grain_model class.) */
  T_unit_map units_;
public:
  scatterer(T_lambda kk, T_lambda aa) :
    opacity_ (independent_copy(kk)), albedo_ (independent_copy (aa)) {
    ASSERT_ALL (opacity_ >= 0); 
    ASSERT_ALL (albedo_ >= 0); 
    ASSERT_ALL (albedo_ <= 1);
 };

  virtual ~scatterer() {};

  /// Returns the opacity vector of the scatterer.
  const T_lambda& opacity() const {return opacity_;};
  /// Returns the albedo vector of the scatterer.
  const T_lambda& albedo() const {return albedo_;};

  /** Returns a lambda vector of 0. Is guaranteed to not be shared, so
      can be copied by reference safely. */
  T_lambda zero_lambda () const {return initialize(opacity (),0.);};
  
  /// Sets the wavelength vector used, which recalculates the opacity
  /// and albedo vectors.
  virtual void set_wavelength(const T_lambda& lambda)=0;
  /// Returns the phase function of the scatterer.
  virtual const angular_distribution<T_lambda>& phase_function () const = 0;

  /** Scatters a ray. Returns the angular distribution object
      describing the angular dependence of the scattered radiation. */
  virtual const angular_distribution<T_lambda>& 
  scatter (T_biaser b, T_ray& r, T_rng_policy& rng, T_float& prob) const=0;

  /// Exception class, thrown if constructor fails.
  class invalid_data {};

  const T_unit_map& units() const {return units_;};
};


/** The HG phase function. It is in the mcrx namespace so that we can
    easily define a blitz expression template version of it. */
inline mcrx::T_float mcrx::HG_phase_function(T_float g, T_float cos_theta) {
  return (1.0-g*g)/((4*constants::pi)*pow (1.0 + g*g - (2*cos_theta)*g, 1.5));
}


/** This class defines the return type of the HG_phase_function. This
    is necessary because for polychromatic situations it returns an
    expression template.  */
template<typename T>
class mcrx::hg_returntype {
public:
  typedef T T_type;
};


/** This class defines the return type of the HG_phase_function. This
    is necessary because for polychromatic situations it returns an
    expression template.  */
template<typename T>
class mcrx::hg_returntype<blitz::Array<T, 1> > {
public:
  typedef typename blitz::BzBinaryExprResult<Fn_hgpf,
					     blitz::Array<T, 1>,
					     T_float>::T_result T_type;
};


/** Dust grain in Henyey-Greenstein scattering approximation. This
 dust grain takes its angular distribution from the HG formula, given
 by opacity, albedo and scattering asymmetry.  */
template <template<class> class chromatic_policy, 
	  typename T_rng_policy = mcrx::global_random> 
class mcrx::simple_HG_dust_grain : 
  public scatterer<chromatic_policy, T_rng_policy> {
public:
  typedef typename scatterer<chromatic_policy, T_rng_policy>::T_lambda 
  T_lambda;
  typedef typename scatterer<chromatic_policy, T_rng_policy>::T_biaser 
  T_biaser;
  typedef typename scatterer<chromatic_policy, T_rng_policy>::T_ray
  T_ray ;

protected:
  /** Angular_distribution object representing the phase function for a
      Henyey-Greenstein dust grain. */
  class hg_phase_function : public angular_distribution<T_lambda>  {
  private:
    T_lambda g_;
  public:
    hg_phase_function (): angular_distribution<T_lambda> (), g_ (0) {};
    hg_phase_function (const T_lambda& gg): 
      angular_distribution<T_lambda> (), g_ (independent_copy(gg)) {};
    void set_g (const T_lambda& gg) {
      assign(g_, gg); 
      ASSERT_ALL (g_ >= -1); 
      ASSERT_ALL (g_ <= 1);};
    const T_lambda& asymmetry () const {return g_;};
    T_lambda distribution_function (const vec3d& new_direction,
				    const vec3d& old_direction) const {
      assert (dot (new_direction, new_direction) > .9999);
      assert (dot (new_direction, new_direction)< 1.0001);
      assert (dot (old_direction, old_direction) > .9999);
      assert (dot (old_direction, old_direction)< 1.0001);
      const T_float cos_theta = dot (new_direction, old_direction);
      return make_thread_local_copy(phase_function(cos_theta));
    };
    void distribution_function (const vec3d& new_direction,
				const vec3d& old_direction,
				T_lambda& distr) const {
      assert (dot (new_direction, new_direction) > .9999);
      assert (dot (new_direction, new_direction)< 1.0001);
      assert (dot (old_direction, old_direction) > .9999);
      assert (dot (old_direction, old_direction)< 1.0001);
      const T_float cos_theta = dot (new_direction, old_direction);
      distr = phase_function(cos_theta);
    };
    typename hg_returntype<T_lambda>::T_type
    phase_function (T_float cos_theta) const {
      return mcrx::HG_phase_function(g_, cos_theta) ;
    };
    boost::shared_ptr<angular_distribution<T_lambda> > clone() const {
      return boost::shared_ptr<angular_distribution<T_lambda> >
	(new hg_phase_function(*this));};
  };

  
  hg_phase_function phfunc;

  /// Default constructor is intended for derived classes.
  simple_HG_dust_grain () {};
public:
  simple_HG_dust_grain (const T_lambda& kk, const T_lambda& gg,
			const T_lambda& aa):
    scatterer<chromatic_policy, T_rng_policy>(kk,aa), phfunc(gg) {};
  // simple_HG_dust_grain (const HG_dust_grain&); //implicitly generated
  //~simple_HG_dust_grain(); //implicitly generated

  const T_lambda& asymmetry () const {return phfunc.asymmetry();};
  
  /// No wavelength dependence, so just a no-op, though it should
  /// probably make sure the vectors are sized appropriately, if
  /// applicable.
  virtual void set_wavelength (const T_lambda& lambda) {};
  
  const angular_distribution<T_lambda>& 
  scatter (T_biaser b, T_ray& r, T_rng_policy& rng, T_float& prob) const;

  const angular_distribution<T_lambda>& phase_function () const {return phfunc;};
};


/** A wavelength-dependent Henyey-Greenstein dust grain that reads its
    properties from a file. */
template <template<class> class chromatic_policy, 
	  typename T_rng_policy = mcrx::global_random> 
class mcrx::HG_dust_grain : 
    public simple_HG_dust_grain<chromatic_policy, T_rng_policy> {
public:
  typedef typename simple_HG_dust_grain<chromatic_policy, T_rng_policy>::T_lambda 
  T_lambda;
  typedef typename simple_HG_dust_grain <chromatic_policy, T_rng_policy>::T_biaser 
  T_biaser;
  typedef typename simple_HG_dust_grain <chromatic_policy, T_rng_policy>::T_ray
  T_ray ;

protected:
  typedef interpolator<T_float,  T_float, 1> T_interpolator ;
  
  // Interpolators containing the dust characteristics
  
  // The use of shared_ptr means that the implicitly generated copy
  // constructor works fine
  boost::shared_ptr<T_interpolator> kappaint;
  boost::shared_ptr<T_interpolator> gint;
  boost::shared_ptr<T_interpolator> aint;

  /// Default constructor is intended for derived classes.
  HG_dust_grain () {}; 

public:
  HG_dust_grain(const std::string&); // Loads dust grain data from file
  //HG_dust_grain (const HG_dust_grain&); //implicitly generated
  //~HG_dust_grain(); //implicitly generated

  // Sets up variables for a certain wavelength 
  virtual void set_wavelength (const T_lambda& lambda);
};

/// The Draine_grain is just a HG_dust_grain that can read the Draine files.
template <template<class> class chromatic_policy, 
	  typename T_rng_policy = mcrx::global_random> 
class mcrx::Draine_grain : 
    public HG_dust_grain<chromatic_policy, T_rng_policy> {
public:
  typedef typename HG_dust_grain<chromatic_policy, T_rng_policy>::T_lambda 
  T_lambda;
  typedef typename HG_dust_grain<chromatic_policy, T_rng_policy>::T_interpolator T_interpolator;

  Draine_grain (const std::string&);
};

 
// **** function definitions *** 


/** Scatters a ray. The ray r is scattered according to the phase
 function of the grain. The analytical formula
 \f$\cos\theta={1+g^2-\left({1-g^2 \over 1-g+2gR_2}\right)^2 \over
 2g}\f$ is used to calculate the new \f$\theta\f$ direction. None of
 the members are affected by the reference wavelength, so there are no
 critical sections here. */
template <template<class> class chromatic_policy, typename T_rng_policy>
const mcrx::angular_distribution<
  typename mcrx::simple_HG_dust_grain<chromatic_policy, 
				      T_rng_policy>::T_lambda>& 
mcrx::simple_HG_dust_grain<chromatic_policy, 
			   T_rng_policy>::scatter (T_biaser b,
						   T_ray& r, 
						   T_rng_policy& rng,
						   T_float& prob) const
{
  // Get the random numbers
  blitz::TinyVector<T_float, 2> rn;
  rng.rnd(rn);

  // Get the new theta from the phase function, and an isotropic phi
  // in a system where the z axis is aligned with the incoming ray
  const T_float phi= rn [0]*2*constants::pi;
  const T_float R2= rn [1];
  // Use the analytical formula to calculate the cos theta.
  const T_float gref = b.reference1(phfunc.asymmetry());
  const T_float costheta= (gref != 0)?
    // remove pow call here, it's unneccessary
  ( (1+gref*gref) - pow((1-gref*gref)/(1-gref+2*gref*R2),2) )/(2*gref) : 
  2*R2- 1;

  // set the probability of the drawn scattering as HG(costheta),
  // which already has the 1/4pi in it.
  prob = b.reference1(phfunc.phase_function(costheta));

  // The ray knows how to scatter itself into a certain direction
  r.scatter_in_direction(costheta, phi);

  // Apply bias to outgoing ray intensity due to the phase function
  // phase function bias is the phase_function(cos_theta)/ref.
  DEBUG(1,T_lambda bias(b.bias_factor(phfunc.phase_function(costheta))); \
	if(!condition_all(bias<1e3)) \
	std::cout << "WARNING: Large bias factor  " << max(bias)<< " at lambda " << maxIndex(bias) << " reflambda " << b.reference_index() << " at " << __FILE__ <<':' <<__LINE__ << std::endl;);

  // note that albedo in not part of this bias, because that's done
  // separately in xfer::scatter.
  r.bias(b.bias_factor(phfunc.phase_function(costheta)));
  ASSERT_ALL(r.intensity()>= 0);

  return phfunc;
}


template <template<class> class chromatic_policy, 
	  typename T_rng_policy> 
mcrx::HG_dust_grain<chromatic_policy, 
		    T_rng_policy>::HG_dust_grain(const std::string& dustname) :
		      simple_HG_dust_grain<chromatic_policy, T_rng_policy> ()
{
  {
    std::ostringstream os;
    os << dustname << ".kappa";
    kappaint = read_old_interpolator<T_float> (os.str ());
  }
  {
    std::ostringstream os;
    os << dustname << ".asym";
    gint = read_old_interpolator<T_float>  (os.str ());
  }
  {
    std::ostringstream os;
    os << dustname << ".albedo";
    aint = read_old_interpolator<T_float>  (os.str ());
  }
}

/** Set the grain parameters at a given wavelength. This function is
    normally called from the dust_model set_wavelength. */
template <template<class> class chromatic_policy, typename T_rng_policy> 
void
mcrx::HG_dust_grain<chromatic_policy, 
		    T_rng_policy>::set_wavelength (const T_lambda& lambda)
{
  // The interpolate should be overloaded for an array lambda so this
  // should be transparent. The opacity array is already resized by
  // the dust model, so we can do a straight assignment here. 

  // we keep mutex locking on here because the opacity array is shared
  // between the dust model and the scatterer and is not guaranteed to
  // be thread local.
  {
    const array_1 old_lambda(reference_vector(kappaint->axis(0)));
    assign(this->opacity_, subsample(kappaint->data(),old_lambda,lambda), true);
  }
  {
    const array_1 old_lambda(reference_vector(aint->axis(0)));
    assign(this->albedo_, subsample(aint->data(),old_lambda,lambda));
  }
  {
    const array_1 old_lambda(reference_vector(gint->axis(0)));
    this->phfunc.set_g(subsample(gint->data(),old_lambda,lambda));
  }
  ASSERT_ALL(this->opacity() >= 0);
  ASSERT_ALL(this->albedo() >= 0);
  ASSERT_ALL(this->albedo() <= 1);
}

template <template<class> class chromatic_policy, typename T_rng_policy> 
mcrx::Draine_grain<chromatic_policy, 
		   T_rng_policy>::Draine_grain (const std::string& file) :
  HG_dust_grain<chromatic_policy, T_rng_policy>  ()
{
  std::ifstream infile (file.c_str ());
  /// \todo DOES NOT WORK FOR D03 FILES BECAUSE OF THE SPURIOUS TEXT CRAP AT
  /// THE ENDLINES!!!

  // First spool up until the numbers start
  bool D03_file = false;
  {
    std::string temp;
    while (getline (infile, temp),
	   // look for cos^2 as a sign that it is a Draine03 file
	   D03_file |= (temp.find("<cos^2>")!= std::string::npos),
	   temp.substr(0, 4) != "----" && infile.good () ) ;
  }

  if (!infile.good ())
    // we did not find the start line
    throw typename scatterer<chromatic_policy, T_rng_policy>::invalid_data ();
  
  // now read the numbers
  std::vector<T_float> input ((std::istream_iterator<T_float> (infile)),
			      std::istream_iterator<T_float> ()) ;

  if ((D03_file && (input.size()% 6)) ||
      (!D03_file && (input.size()%5)) ||
      (input.size() == 0))
    // table has five columns (Draine 03 has 6), if not, it's not the
    // correct file
    throw typename scatterer<chromatic_policy, T_rng_policy>::invalid_data ();

  // and calculate the values
  std::vector<T_float> lambda, g, albedo, k ;
  for (std::vector<T_float>::iterator i = input.begin() ; i < input.end() ; ) {
    lambda.push_back(*i++ *1e-6); // wavelength is in Micron
    assert (i != input.end());
    albedo.push_back(*i++);
    g.push_back(*i++);
    // need to skip C_ext in the file
    // 2.089e-10 is the conversion from cm^2/g to kpc^2/M_Sun
    k.push_back(*((++i)++) *2.089e-10/(1 -albedo.back())); 
    // since were currently not using cos^2, we skip that too
    if (D03_file) ++i;
    //std::cout << lambda.back() << ' ' << g.back() << ' ' << albedo.back() << std::endl;
  }
  
  if(D03_file) {
    // The D03 files are in order of DECREASING wavelength... retarded
    reverse(lambda.begin(), lambda.end());
    reverse(g.begin(), g.end());
    reverse(albedo.begin(), albedo.end());
    reverse(k.begin(), k.end());
  }

  this->kappaint.reset(new T_interpolator) ;
  this->kappaint->setAxis(0 , lambda) ;
  this->gint.reset(new T_interpolator);
  this->gint->setAxis(0 , lambda) ;
  this->aint.reset(new T_interpolator)  ;
  this->aint->setAxis(0 , lambda) ;

  for (int i = 0; i < lambda.size() ; ++i) {
    this->kappaint->setPoint(i, k [i]);
    this->gint->setPoint(i, g [i]);
    this->aint->setPoint(i, albedo [i]);
  } 

}


#endif


